/*******************************************************************************
 * Copyright (c) 2025-08-06 @author <a href="mailto:iffiff1@gmail.com">Tyler Chen</a>.
 * All rights reserved.
 *
 * Contributors:
 *     <a href="mailto:iffiff1@gmail.com">Tyler Chen</a> - initial API and implementation.
 ******************************************************************************/
package org.iff.util.llm.qwen;

import org.apache.commons.lang3.StringUtils;
import org.iff.util.HttpClientUtil;
import org.iff.util.JSONUtil;
import org.iff.util.Utils;

import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;

/**
 * Qwen
 *
 * @author <a href="mailto:iffiff1@gmail.com">Tyler Chen</a>
 * @since 2025-08-06
 */
public class Qwen {
    public static final String URI = "https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions";

    public static String chat(String apiKey, String uri, String model, String system, String user) {
        try {
            String json = JSONUtil.toJsonString(TextMessage.of()
                    .stream(false)
                    .model(model)
                    .message(TextMessage.ROLE_SYSTEM, system)
                    .message(TextMessage.ROLE_USER, user));
            HttpClientUtil.HttpOptions httpOptions = HttpClientUtil.HttpOptions.create().authorization("Bearer " + apiKey)
                    .header("Content-Type", "application/json; charset=UTF-8");
            String post = HttpClientUtil.post(uri, json, httpOptions);
            TextChat textChat = JSONUtil.fromJson(cleanJson(post), TextChat.class);
            if (textChat != null && textChat.getChoices() != null) {
                List<String> collect = Utils.collect(textChat.getChoices(), c -> c.getMessage() == null ? "" : c.getMessage().getContent());
                return StringUtils.join(collect, '\n');
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String chatStream(String apiKey, String uri, String model, String system, String user) {
        try {
            String json = JSONUtil.toJsonString(TextMessage.of()
                    .stream(true)
                    .model(model)
                    .message(TextMessage.ROLE_SYSTEM, system)
                    .message(TextMessage.ROLE_USER, user));
            List<TextChatStream> result = new ArrayList<>();
            Consumer<String> receive = (line) -> {
                if (line != null && line.length() > 0) {
                    try {
                        if (line.indexOf('{') >= 0) {
                            TextChatStream textChatStream = JSONUtil.fromJson(cleanJson(line), TextChatStream.class);
                            if (textChatStream != null) {
                                result.add(textChatStream);
                            }
                        }
                    } catch (Exception e) {
                        System.out.println("Error content: " + line);
                        e.printStackTrace();
                    }
                }
            };
            HttpClientUtil.HttpOptions httpOptions = HttpClientUtil.HttpOptions.create().authorization("Bearer " + apiKey)
                    .header("Content-Type", "application/json; charset=UTF-8");
            HttpClientUtil.postAcceptLine(uri, json.getBytes(StandardCharsets.UTF_8), httpOptions, receive, null);
            StringBuilder sb = new StringBuilder(128);
            for (TextChatStream textChatStream : result) {
                if (textChatStream.getChoices() != null) {
                    for (TextChatStream.Choices choice : textChatStream.getChoices()) {
                        if (choice.getDelta() != null && choice.getDelta().getContent() != null) {
                            sb.append(choice.getDelta().getContent());
                        }
                    }
                }
            }
            return sb.toString();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String chatThink(String apiKey, String uri, String model, String system, String user) {
        try {
            String json = JSONUtil.toJsonString(TextMessage.of()
                    .stream(true).enableThinking(true)
                    .model(model)
                    .message(TextMessage.ROLE_SYSTEM, system)
                    .message(TextMessage.ROLE_USER, user));
            List<TextChatStream> result = new ArrayList<>();
            Consumer<String> receive = (line) -> {
                if (line != null && line.length() > 0) {
                    try {
                        if (line.indexOf('{') >= 0) {
                            TextChatStream textChatStream = JSONUtil.fromJson(cleanJson(line), TextChatStream.class);
                            if (textChatStream != null) {
                                result.add(textChatStream);
                            }
                        }
                    } catch (Exception e) {
                        System.out.println("Error content: " + line);
                        e.printStackTrace();
                    }
                }
            };
            HttpClientUtil.HttpOptions httpOptions = HttpClientUtil.HttpOptions.create().authorization("Bearer " + apiKey)
                    .header("Content-Type", "application/json; charset=UTF-8");
            HttpClientUtil.postAcceptLine(uri, json.getBytes(StandardCharsets.UTF_8), httpOptions, receive, null);
            StringBuilder think = new StringBuilder(128).append("<think>\n");
            StringBuilder content = new StringBuilder(128);
            for (TextChatStream textChatStream : result) {
                if (textChatStream.getChoices() != null) {
                    for (TextChatStream.Choices choice : textChatStream.getChoices()) {
                        if (choice.getDelta() != null && choice.getDelta().getContent() != null) {
                            content.append(choice.getDelta().getContent());
                        } else if (choice.getDelta() != null && choice.getDelta().getReasoningContent() != null) {
                            think.append(choice.getDelta().getReasoningContent());
                        }
                    }
                }
            }
            return think.append("\n</think>\n").append(content).toString();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    static String cleanJson(String json) {
        if (json == null) {
            return json;
        }
        int start = json.indexOf('{');
        int end = json.lastIndexOf('}');
        if (start >= 0 && end > start) {
            return json.substring(start, end + 1);
        }
        return json;
    }
}
